{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08327e3c",
   "metadata": {},
   "source": [
    "# dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b619eabd",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b9bcc48d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cannot find config for target=ext_dev -keys=vta,cpu -device=vta -model=sim_1x16_i8w8a32_15_15_18_17, workload=('dense_packed.vta', ('TENSOR', (16, 32, 1, 16), 'int8'), ('TENSOR', (63, 32, 16, 16), 'int8'), None, 'int32'). A fallback configuration is used, which may bring great performance regression.\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.uop_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.command_handle\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_push\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_sync\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/script/printer/tir/expr.cc:246: Warning: No TScriptPrinterName attribute for tir.vta.coproc_dep_pop\n",
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# from tvm.script import ir as I\n",
      "# from tvm.script import tir as T\n",
      "\n",
      "@I.ir_module\n",
      "class Module:\n",
      "    @T.prim_func\n",
      "    def main(data: T.Buffer((16, 32, 1, 16), \"int8\"), kernel: T.Buffer((63, 32, 16, 16), \"int8\"), compute: T.Buffer((16, 63, 1, 16), \"int8\")):\n",
      "        T.func_attr({\"from_legacy_te_schedule\": T.bool(True), \"tir.noalias\": T.bool(True)})\n",
      "        T.tir.vta.coproc_dep_push(3, 2)\n",
      "        for i0_outer, i1_outer in T.grid(16, 63):\n",
      "            vta = T.int32()\n",
      "            with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_scope\", 2):\n",
      "                T.tir.vta.coproc_dep_pop(3, 2)\n",
      "                with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_uop_scope\", \"VTAPushGEMMOp\"):\n",
      "                    T.tir.vta.uop_push(0, 1, 0, 0, 0, 0, 0, 0)\n",
      "                T.tir.vta.coproc_dep_push(2, 1)\n",
      "            for k_o_outer in range(32):\n",
      "                with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_scope\", 1):\n",
      "                    T.tir.vta.coproc_dep_pop(2, 1)\n",
      "                    T.call_extern(\"int32\", \"VTALoadBuffer2D\", T.tvm_thread_context(T.tir.vta.command_handle()), data.data, i0_outer * 32 + k_o_outer, 1, 1, 1, 0, 0, 0, 0, 0, 2)\n",
      "                    T.call_extern(\"int32\", \"VTALoadBuffer2D\", T.tvm_thread_context(T.tir.vta.command_handle()), kernel.data, i1_outer * 32 + k_o_outer, 1, 1, 1, 0, 0, 0, 0, 0, 1)\n",
      "                    T.tir.vta.coproc_dep_push(1, 2)\n",
      "                T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_scope\", 2)\n",
      "                T.tir.vta.coproc_dep_pop(1, 2)\n",
      "                with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_uop_scope\", \"VTAPushGEMMOp\"):\n",
      "                    T.tir.vta.uop_push(0, 0, 0, 0, 0, 0, 0, 0)\n",
      "                T.tir.vta.coproc_dep_push(2, 1)\n",
      "            T.tir.vta.coproc_dep_pop(2, 1)\n",
      "            with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_scope\", 2):\n",
      "                with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_uop_scope\", \"VTAPushALUOp\"):\n",
      "                    T.tir.vta.uop_push(1, 0, 0, 0, 0, 3, 1, 8)\n",
      "                with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_uop_scope\", \"VTAPushALUOp\"):\n",
      "                    T.tir.vta.uop_push(1, 0, 0, 0, 0, 0, 1, 127)\n",
      "                with T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_uop_scope\", \"VTAPushALUOp\"):\n",
      "                    T.tir.vta.uop_push(1, 0, 0, 0, 0, 1, 1, 0)\n",
      "                T.tir.vta.coproc_dep_push(2, 3)\n",
      "            T.attr(T.iter_var(vta, None, \"ThreadIndex\", \"vta\"), \"coproc_scope\", 3)\n",
      "            T.tir.vta.coproc_dep_pop(2, 3)\n",
      "            T.call_extern(\"int32\", \"VTAStoreBuffer2D\", T.tvm_thread_context(T.tir.vta.command_handle()), 0, 4, compute.data, i0_outer * 63 + i1_outer, 1, 1, 1)\n",
      "            T.tir.vta.coproc_dep_push(3, 2)\n",
      "        T.tir.vta.coproc_sync()\n",
      "        T.tir.vta.coproc_dep_pop(3, 2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:03:03.969 INFO load_module /tmp/tmpox3977y4/dense.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA DENSE TEST PASSED: Time cost = 0.0187377 sec/op, 0.88138 GOPS\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:03:03] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/runtime/profiling.cc:101: Warning: No timer implementation for ext_dev, using default timer instead. It may be inaccurate or have extra overhead.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "from collections import namedtuple\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import tvm\n",
    "from tvm import te\n",
    "from tvm import autotvm\n",
    "from tvm.contrib import utils\n",
    "from tvm.contrib.pickle_memoize import memoize\n",
    "from tvm import topi\n",
    "import tvm.topi.testing\n",
    "import vta\n",
    "from vta import program_fpga, reconfig_runtime\n",
    "import vta.testing\n",
    "from vta.testing import simulator\n",
    "\n",
    "# FIXME: we need a custom clip operator to circumvent a pattern detection limitation\n",
    "@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)\n",
    "def my_clip(x, a_min, a_max):\n",
    "    \"\"\"Unlike topi's current clip, put min and max into two stages.\"\"\"\n",
    "    const_min = tvm.tir.const(a_min, x.dtype)\n",
    "    const_max = tvm.tir.const(a_max, x.dtype)\n",
    "    x = te.compute(x.shape, lambda *i: tvm.te.min(x(*i), const_max), name=\"clipA\")\n",
    "    x = te.compute(x.shape, lambda *i: tvm.te.max(x(*i), const_min), name=\"clipB\")\n",
    "    return x\n",
    "\n",
    "\n",
    "def run_gemm(\n",
    "    env,\n",
    "    remote,\n",
    "    target,\n",
    "    batch_size,\n",
    "    in_feat,\n",
    "    out_feat,\n",
    "    check_correctness=True,\n",
    "    print_ir=True,\n",
    "    samples=4,\n",
    "):\n",
    "\n",
    "    # Perform packing only if we are targeting the accelerator\n",
    "    if \"arm_cpu\" in target.keys:\n",
    "        data_pack = False\n",
    "    elif \"vta\" in target.keys:\n",
    "        data_pack = True\n",
    "\n",
    "    # Derive shapes depending upon packing\n",
    "    a_shape = (batch_size, in_feat)\n",
    "    w_shape = (out_feat, in_feat)\n",
    "    if data_pack:\n",
    "        data_shape = (batch_size // env.BATCH, in_feat // env.BLOCK_IN, env.BATCH, env.BLOCK_IN)\n",
    "        kernel_shape = (\n",
    "            out_feat // env.BLOCK_OUT,\n",
    "            in_feat // env.BLOCK_IN,\n",
    "            env.BLOCK_OUT,\n",
    "            env.BLOCK_IN,\n",
    "        )\n",
    "        fcompute = vta.top.dense_packed\n",
    "        fschedule = vta.top.schedule_dense_packed\n",
    "    else:\n",
    "        data_shape = a_shape\n",
    "        kernel_shape = w_shape\n",
    "        fcompute = topi.x86.dense_nopack\n",
    "        fschedule = topi.x86.schedule_dense_nopack\n",
    "    data = te.placeholder(data_shape, name=\"data\", dtype=env.inp_dtype)\n",
    "    kernel = te.placeholder(kernel_shape, name=\"kernel\", dtype=env.wgt_dtype)\n",
    "\n",
    "    # Define base computation schedule\n",
    "    with target:\n",
    "        res = fcompute(data, kernel, None, env.acc_dtype)\n",
    "        res = topi.right_shift(res, 8)\n",
    "        res = my_clip(res, 0, (1 << env.OUT_WIDTH - 1) - 1)\n",
    "        res = topi.cast(res, env.out_dtype)\n",
    "        # Derive base schedule\n",
    "        s = fschedule([res])\n",
    "        if print_ir:\n",
    "            print(vta.lower(s, [data, kernel, res], simple_mode=True))\n",
    "\n",
    "    # Derive number of ops\n",
    "    num_ops = 2 * batch_size * in_feat * out_feat\n",
    "\n",
    "    # @memoize(\"vta.tests.test_benchmark_topi.dense.verify\")\n",
    "    def get_ref_data():\n",
    "        # derive min max for act, wgt types (max non inclusive)\n",
    "        a_min, a_max = 0 - (1 << (env.INP_WIDTH - 1)), (1 << (env.INP_WIDTH - 1))\n",
    "        w_min, w_max = 0 - (1 << (env.WGT_WIDTH - 1)), (1 << (env.WGT_WIDTH - 1))\n",
    "        a_np = np.random.randint(a_min, a_max, size=a_shape).astype(data.dtype)\n",
    "        w_np = np.random.randint(w_min, w_max, size=w_shape).astype(kernel.dtype)\n",
    "\n",
    "        r_np = np.dot(a_np.astype(env.acc_dtype), w_np.T.astype(env.acc_dtype)).astype(\n",
    "            env.acc_dtype\n",
    "        )\n",
    "        return a_np, w_np, r_np\n",
    "\n",
    "    # Data in original format\n",
    "    data_np, kernel_np, res_ref = get_ref_data()\n",
    "    if data_pack:\n",
    "        data_np = data_np.reshape(\n",
    "            batch_size // env.BATCH, env.BATCH, in_feat // env.BLOCK_IN, env.BLOCK_IN\n",
    "        ).transpose((0, 2, 1, 3))\n",
    "        kernel_np = kernel_np.reshape(\n",
    "            out_feat // env.BLOCK_OUT, env.BLOCK_OUT, in_feat // env.BLOCK_IN, env.BLOCK_IN\n",
    "        ).transpose((0, 2, 1, 3))\n",
    "\n",
    "    # Build\n",
    "    if \"vta\" in target.keys:\n",
    "        mod = vta.build(\n",
    "            s,\n",
    "            [data, kernel, res],\n",
    "            target=tvm.target.Target(target, host=env.target_host),\n",
    "            name=\"dense\",\n",
    "        )\n",
    "    else:\n",
    "        mod = tvm.build(\n",
    "            s,\n",
    "            [data, kernel, res],\n",
    "            target=tvm.target.Target(target, host=env.target_host),\n",
    "            name=\"dense\",\n",
    "        )\n",
    "    temp = utils.tempdir()\n",
    "    mod.save(temp.relpath(\"dense.o\"))\n",
    "    remote.upload(temp.relpath(\"dense.o\"))\n",
    "    f = remote.load_module(\"dense.o\")\n",
    "    dev = remote.device(str(target))\n",
    "\n",
    "    res_np = np.zeros(topi.utils.get_const_tuple(res.shape)).astype(res.dtype)\n",
    "    data_arr = tvm.nd.array(data_np, dev)\n",
    "    kernel_arr = tvm.nd.array(kernel_np, dev)\n",
    "    res_arr = tvm.nd.array(res_np, dev)\n",
    "    time_f = f.time_evaluator(\"dense\", dev, number=samples)\n",
    "\n",
    "    # In vta sim mode, collect simulator runtime statistics\n",
    "    stats = {}\n",
    "    cost = None\n",
    "    if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "        # Check if we're in local RPC mode (allows us to rebuild the\n",
    "        # runtime on the fly when varying the VTA designs)\n",
    "        local_rpc = int(os.environ.get(\"VTA_LOCAL_SIM_RPC\", \"0\"))\n",
    "        if local_rpc:\n",
    "            if env.TARGET == \"sim\":\n",
    "                remote.get_function(\"vta.simulator.profiler_clear\")()\n",
    "            else:\n",
    "                remote.get_function(\"vta.tsim.profiler_clear\")()\n",
    "            cost = time_f(data_arr, kernel_arr, res_arr)\n",
    "            if env.TARGET == \"sim\":\n",
    "                stats = json.loads(remote.get_function(\"vta.simulator.profiler_status\")())\n",
    "            else:\n",
    "                stats = json.loads(remote.get_function(\"vta.tsim.profiler_status\")())\n",
    "        else:\n",
    "            simulator.clear_stats()\n",
    "            cost = time_f(data_arr, kernel_arr, res_arr)\n",
    "            stats = simulator.stats()\n",
    "    else:\n",
    "        cost = time_f(data_arr, kernel_arr, res_arr)\n",
    "\n",
    "    # Check correctness\n",
    "    correct = False\n",
    "    if check_correctness:\n",
    "        res_orig = res_arr.numpy()\n",
    "        if data_pack:\n",
    "            res_orig = res_orig.reshape(batch_size, out_feat)\n",
    "        res_ref = res_ref >> 8\n",
    "        res_ref = np.clip(res_ref, 0, (1 << env.OUT_WIDTH - 1) - 1)\n",
    "        res_ref = res_ref.astype(env.out_dtype)\n",
    "        correct = np.allclose(res_orig, res_ref)\n",
    "\n",
    "    gops = (num_ops / cost.mean) / float(10**9)\n",
    "    status = \"PASSED\" if correct else \"FAILED\"\n",
    "    if \"arm_cpu\" in target.keys:\n",
    "        device = \"CPU\"\n",
    "    elif \"vta\" in target.keys:\n",
    "        device = \"VTA\"\n",
    "    print(\"%s DENSE TEST %s: Time cost = %g sec/op, %g GOPS\" % (device, status, cost.mean, gops))\n",
    "\n",
    "    return correct, cost, stats\n",
    "\n",
    "\n",
    "def test_gemm(device=\"vta\", batch=128, in_feat=128, out_feat=128):\n",
    "    def _run(env, remote):\n",
    "        if device == \"vta\":\n",
    "            target = env.target\n",
    "            if env.TARGET not in [\"sim\", \"tsim\"]:\n",
    "                assert tvm.runtime.enabled(\"rpc\")\n",
    "                program_fpga(remote, bitstream=None)\n",
    "                reconfig_runtime(remote)\n",
    "        elif device == \"arm_cpu\":\n",
    "            target = env.target_vta_cpu\n",
    "        with autotvm.tophub.context(target):  # load pre-tuned schedule parameters\n",
    "            run_gemm(env, remote, target, batch, in_feat, out_feat)\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_gemm(\"vta\", 16, 512, 1008)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
